package com.tianji.aigc.controller;

import cn.hutool.core.collection.CollStreamUtil;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;

import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Slf4j
@RestController
@RequestMapping("/embedding")
@RequiredArgsConstructor
public class EmbeddingController {

    private final VectorStore vectorStore;
    private final EmbeddingModel embeddingModel;

    //保存到向量数据库
    @PostMapping
    public void saveVectorStore(@RequestParam("messages")List<String> messages){
        log.info("保存到向量数据库中，消息数据：{}", messages);
        //构建文档
        List<Document> documents = CollStreamUtil.toList(messages, message -> Document.builder()
                .text(message)
                .build());

        //保存到向量库中
        vectorStore.add(documents);
        log.info("保存到向量数据库成功, 数量：{}", messages.size());
    }

    //文本转换向量
    @GetMapping
    public EmbeddingResponse textToVector(@RequestParam("message") String message){
        log.info("将文本转换成向量，消息数据：{}", message);
        return embeddingModel.embedForResponse(List.of(message));
    }
    //内容搜索
    @GetMapping("/search")
    public List<Document> search(@RequestParam("message") String message ){
        return vectorStore.similaritySearch(SearchRequest.builder().query(message).topK(5).build());
    }

    //搜索全部
    @GetMapping("/search/all")
    public  List<Document> searchAll(){
        return vectorStore.similaritySearch(SearchRequest.builder().query("").topK(999).build());

    }

    //删除向量数据库中的数据
    @DeleteMapping
    public void deleteText(@RequestParam("ids") List<String> ids){
        vectorStore.delete(ids);
    }

}
